'''
https://leetcode.cn/problems/height-of-binary-tree-after-subtree-removal-queries/description/
'''
from collections import defaultdict
from typing import Optional, List


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:

    def treeQueries(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        dfn = defaultdict(int)
        sizes = defaultdict(int)
        deeps = defaultdict(int)
        id = 1
        def f(root, deep):
            if not root:
                return 0
            nonlocal id
            dfn[root.val] = id
            deeps[id] = deep
            id += 1
            size = f(root.left, deep+1) + f(root.right, deep+1) + 1
            # size[root.val] = size              # 原始下标，统计dfn和size
            sizes[dfn[root.val]] = size      # 原始下标统计dfn，dfn下标统计size
            return size
        f(root, 0)

        # for node in dfn.keys():
        #     dfnid = dfn[node]
        #     print(f"node: {node}, dfnid: {dfnid}, size: {sizes[dfnid]}, deep: {deeps[dfnid]}")
        n = len(dfn)
        maxHeightL, maxHeightR = [0] * (n + 2), [0] * (n + 2)   # 根据dfn编号，左边最大高度，右边最大高度
        for id in range(1, n + 1):
            maxHeightL[id] = max(maxHeightL[id-1], deeps[id])
        maxHeightL[n+1] = maxHeightL[n]
        for id in range(n, 0, -1):
            maxHeightR[id] = max(maxHeightR[id+1], deeps[id])
        maxHeightR[0] = maxHeightR[1]

        res = queries[:]
        for i, query in enumerate(queries):
            dfnid = dfn[query]
            subtreeSize = sizes[dfnid]
            boardL, boardR = dfnid, dfnid + subtreeSize - 1     # 这段区间内的节点都是被删除的
            res[i] = max(maxHeightL[boardL-1], maxHeightR[boardR + 1])
        return res

    def treeQueries2(self, root: Optional[TreeNode], queries: List[int]) -> List[int]:
        # f函数：统计二叉树大小
        def f(root):
            return 1 + f(root.left) + f(root.right) if root else 0
        n = f(root)
        dfn = [0] * (n + 1)
        sizes = [0] * (n + 1)
        deeps = [0] * (n + 1)

        # g函数：填充dfn编号，和统计sizes以dfn编号为下标
        def g(root, id):
            if not root: return
            dfn[root.val] = id      # 先序统计dfn编号
            l_size = g(root.left, id + 1)
            r_size = g(root.right, id + l_size + 1)
            size = 1 + l_size + r_size
            sizes[id] = size        # 后序统计size，以dfn编号为下标
            return size
        # h函数：统计deep以dfn编号为下标
        def h(root, deep):
            if not root: return
            h(root.left, deep+1)
            h(root.right, deep+1)
            deeps[dfn[root.val]] = deep     # 后序统计deep，以dfn编号为下标

        g(root, 1)
        h(root, 0)  # 题目的高度只算边，即可以理解为不包括根节点

        for node in range(1, n+1):
            dfnid = dfn[node]
            print(f"node: {node}, dfnid: {dfnid}, size: {sizes[dfnid]}, deep: {deeps[dfnid]}")
        maxHeightL, maxHeightR = [0] * (n + 2), [0] * (n + 2)   # 根据dfn编号，左边最大高度，右边最大高度
        for dfnid in range(1, n + 1):
            maxHeightL[dfnid] = max(maxHeightL[dfnid-1], deeps[dfnid])
        maxHeightL[n+1] = maxHeightL[n]
        for dfnid in range(n, 0, -1):
            maxHeightR[dfnid] = max(maxHeightR[dfnid+1], deeps[dfnid])
        maxHeightR[0] = maxHeightR[1]

        res = queries[:]
        for i, query in enumerate(queries):
            dfnid = dfn[query]
            subtreeSize = sizes[dfnid]
            boardL, boardR = dfnid, dfnid + subtreeSize - 1     # 这段区间内的节点都是被删除的
            res[i] = max(maxHeightL[boardL-1], maxHeightR[boardR + 1])
        return res

root = TreeNode(1)
root.left = TreeNode(3)
root.right = TreeNode(4)
root.left.left = TreeNode(2)
root.right.left = TreeNode(6)
root.right.right = TreeNode(5)
root.right.right.right = TreeNode(7)
queries = [4, 5, 6]
print(Solution().treeQueries(root, queries))